from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

#分割文档
pdf_path = "./docs/广州城市理工学院合格评估.pdf"

from langchain.schema import Document
from pdf2image import convert_from_path
from PIL import Image, ImageEnhance
import os
import fitz
import io
from langchain.text_splitter import RecursiveCharacterTextSplitter
import easyocr
# 初始化EasyOCR reader
# 提示：首次运行时会自动下载模型（需要网络），请保持耐心
# 如果你只需要中文，可以只保留 'ch_sim'
try:
    # 创建reader实例，指定语言（中文简体、英文）
    reader = easyocr.Reader(['ch_sim'], gpu=True)  # gpu=True 如果你有GPU且配置好了CUDA，可以加速
    print("EasyOCR reader 初始化成功。")
except Exception as e:
    print(f"初始化 EasyOCR 失败: {e}")
    # 根据你的需求，决定是否退出或采用其他处理方式
    # exit(1)


class pdf_parser:
    def __init__(self):
        self.pdf_path = pdf_path
    
    def enhance_image_for_ocr(self,image):
        """
        增强图像以提高OCR准确性
        
        参数:
        image: PIL图像对象
        
        返回:
        PIL图像对象: 增强后的图像
        """
        # 转换为灰度
        image = image.convert('L')
        
        # 增强对比度
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(2.0)
        
        # 增强锐度
        enhancer = ImageEnhance.Sharpness(image)
        image = enhancer.enhance(2.0)
        
        return image

    def extract_text_from_pdf_images(self, min_text_length=10):
        """
        从PDF中的图像提取文本（智能版）：只对包含图像的页面进行OCR
        
        参数:
        pdf_path (str): PDF文件的路径
        min_text_length (int): 判断页面是否为图像页的最小文本长度阈值
        
        返回:
        list: 包含提取文本的Document对象列表
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(self.pdf_path):
                raise FileNotFoundError(f"文件不存在: {self.pdf_path}")
            
            # 打开PDF文件
            doc = fitz.open(self.pdf_path)
            extracted_docs = []
            
            print(f"开始处理PDF: {self.pdf_path}")
            
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                
                # 首先尝试提取文本
                text = page.get_text()
                
                # 如果提取的文本长度小于阈值，则认为可能是图像页面
                if len(text.strip()) < min_text_length:
                    print(f"第 {page_num+1} 页可能是图像页面，尝试OCR...")
                    
                    # 获取页面中的所有图像
                    image_list = page.get_images()
                    
                    if image_list:
                        print(f"在第 {page_num+1} 页找到 {len(image_list)} 张图像")
                        
                        # 提取页面中的所有文本（包括图像中的文本通过OCR）
                        full_text = ""
                        
                        # 将整个页面渲染为图像进行OCR（适用于扫描件）
                        mat = fitz.Matrix(2, 2)  # 高分辨率渲染
                        pix = page.get_pixmap(matrix=mat)
                        img_data = pix.tobytes("ppm")
                        
                        # 转换为PIL图像
                        image = Image.open(io.BytesIO(img_data))
                        
                        # 增强图像以提高OCR准确性
                        enhanced_image = self.enhance_image_for_ocr(image) 
                        # 注意：easyocr 需要numpy数组格式的图像
                        import numpy as np
                        enhanced_image_np = np.array(enhanced_image)
                        # 使用pytesseract进行OCR
                        # ocr_text = pytesseract.image_to_string(
                        #     enhanced_image, 
                        #     lang='chi_sim+eng',  # 中文+英文
                        #     config='--psm 6'  # 假设为统一的文本块
                        # )
                        # 使用 easyocr 进行OCR识别
                        # result 是一个列表，每个元素是 [坐标, 文本]
                        ocr_results = reader.readtext(enhanced_image_np, paragraph=True) # paragraph=True 尝试组合段落
                        # print("our OCR results:",ocr_results[0])
                        # 提取识别出的文本
                        for _, ocr_text in ocr_results:
                            full_text += ocr_text + "\n"
                        # 如果提取到了文本，创建Document对象
                        if full_text.strip():
                            doc_obj = Document(
                                page_content=full_text,
                                metadata={
                                    "source": pdf_path,
                                    "page": page_num+1,
                                    "type": "ocr_extracted",
                                    "image_count": len(image_list)
                                }
                            )
                            extracted_docs.append(doc_obj)
                            print(f"第 {page_num+1} 页OCR提取成功")
                        else:
                            print(f"第 {page_num+1} 页OCR未提取到文本")
                    else:
                        print(f"第 {page_num+1} 页文本较少但未找到图像，可能只是空白页")
                else:
                    # 对于文本页面，直接使用提取的文本
                    doc_obj = Document(
                        page_content=text,
                        metadata={
                            "source": pdf_path,
                            "page": page_num+1,
                            "type": "text_extracted"
                        }
                    )
                    extracted_docs.append(doc_obj)
                    print(f"第 {page_num+1} 页为文本页面，直接提取文本")
            
            doc.close()
            print(f"处理完成。成功提取了 {len(extracted_docs)} 页文本")
            return extracted_docs
            
        except Exception as e:
            print(f"处理PDF时出错: {str(e)}")
            return []
        
    def pdf_split(self):   
        extracted_docs = self.extract_text_from_pdf_images()
        all_splits = []
        if extracted_docs:
            # 分块处理
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
            all_splits = text_splitter.split_documents(extracted_docs)
            print(f"成功提取并分割了 {len(all_splits)} 个文本块")
        else:
            print("未能从PDF中提取任何文本")
        return all_splits

